package tree

import "fmt"

type Node[V comparable] struct {
	data        V
	left, right *Node[V]
}

type BinaryTree[V comparable] struct {
	root *Node[V]
	cmp  func(V, V) int
}

func NewBinaryTree[V comparable](cmp func(V, V) int) *BinaryTree[V] {
	return &BinaryTree[V]{
		cmp: cmp,
	}
}

func (b *BinaryTree[V]) InsertItem(val V) {
	if b.root == nil {
		b.root = &Node[V]{data: val}
		return
	}

	current := b.root
	for {
		com := b.cmp(val, current.data)
		if com > 0 {
			if current.right == nil {
				current.right = &Node[V]{
					data: val,
				}
				return
			}
			current = current.right
		} else if com < 0 {
			if current.left == nil {
				current.left = &Node[V]{
					data: val,
				}
				return
			}
		} else {
			current.data = val
		}
	}
}

func (b *BinaryTree[V]) ModifyItem(val V) (*Node[V], bool) {
	if b.root == nil {
		return nil, false
	}
	current := b.root
	for current != nil {
		com := b.cmp(val, current.data)
		if com == 0 {
			return current, true
		} else if com > 0 {
			current = current.right
		} else if com < 0 {
			current = current.left
		}
	}
	return nil, false
}

func (b *BinaryTree[V]) delMinNode(node *Node[V]) *Node[V] {
	minNode := node
	current := node
	for {
		if current.left == nil {
			minNode.left = current.right
			break
		}
		current = current.left
		minNode = current
	}
	return minNode
}

func (b *BinaryTree[V]) RemoveItem(val V) bool {
	if b.root == nil {
		return false
	}
	current := b.root
	par := b.root
	for current != nil {

		com := b.cmp(val, current.data)
		if com == 0 {
			//是否根节点
			if current == b.root {
				if current.right == nil {
					b.root = current.left
				} else {
					minNode := b.delMinNode(current.right)
					minNode.left = current.left
					if b.cmp(minNode.data, current.right.data) == 0 {
						minNode.right = nil
					} else {
						minNode.right = current.right
					}
					b.root = minNode
					return true
				}
			}

			if current.left == nil {
				if b.cmp(par.data, current.data) > 0 {
					par.left = current.right
				} else {
					par.right = current.right
				}
				return true
			}
			if current.right == nil {
				if b.cmp(par.data, current.data) > 0 {
					par.left = current.left
				} else {
					par.right = current.left
				}
				return true
			}
			minNode := b.delMinNode(current.right)
			minNode.left = current.left
			if b.cmp(minNode.data, current.data) == 0 {
				minNode.right = nil
			} else {
				minNode.right = current.right
			}
			return true
		} else if com > 0 {
			par = current
			current = current.right
		} else if com < 0 {
			par = current
			current = current.left
		}
	}
	return false
}

func (b *BinaryTree[V]) Preorder() {
	stack := []*Node[V]{b.root}
	for len(stack) > 0 {

		pop := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		fmt.Println(pop.data)

		if pop.right != nil {
			stack = append(stack, pop.right)
		}
		if pop.left != nil {
			stack = append(stack, pop.left)
		}

	}
}
